import os
import concurrent.futures
import random

def process_and_split_txt_file(input_file_path, output_file_paths, ratios, empty_ratios):
    with open(input_file_path, 'r') as f:
        jpg_paths = [line.strip() for line in f.readlines()]

    # 读取所有相关txt文件的内容
    txt_contents = {}
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_jpg_path = {executor.submit(read_txt_file_content, jpg_path): jpg_path for jpg_path in jpg_paths}
        for future in concurrent.futures.as_completed(future_to_jpg_path):
            jpg_path = future_to_jpg_path[future]
            txt_contents[jpg_path] = future.result()

    # 分离有内容和无内容的jpg路径
    empty_txt_jpg_paths = [jpg for jpg, content in txt_contents.items() if not content]
    non_empty_txt_jpg_paths = [jpg for jpg, content in txt_contents.items() if content]

    # 初始化输出列表
    outputs = [[] for _ in output_file_paths]

    # 随机打乱有内容的图片路径
    random.shuffle(non_empty_txt_jpg_paths)

    # 随机打乱无内容的图片路径
    random.shuffle(empty_txt_jpg_paths)

    # 分配有内容图片
    start_index = 0
    for i, ratio in enumerate(ratios):
        if i == len(ratios) - 1:
            # 最后一个文件包含所有剩余的有标注框图片
            outputs[i] = non_empty_txt_jpg_paths[start_index:]
        else:
            end_index = start_index + int(len(non_empty_txt_jpg_paths) * ratio)
            outputs[i] = non_empty_txt_jpg_paths[start_index:end_index]
            start_index = end_index

    # 分配无内容图片
    empty_start_index = 0
    for i, empty_ratio in enumerate(empty_ratios):
        empty_end_index = empty_start_index + int(len(empty_txt_jpg_paths) * empty_ratio)
        outputs[i] += empty_txt_jpg_paths[empty_start_index:empty_end_index]
        empty_start_index = empty_end_index

    # 确保最后一个文件包含所有有标注框图片
    outputs[-1] = non_empty_txt_jpg_paths + outputs[-1]

    # 写入文件
    with concurrent.futures.ThreadPoolExecutor() as executor:
        for file_path, paths in zip(output_file_paths, outputs):
            executor.submit(write_output_file, file_path, paths)

def read_txt_file_content(jpg_path):
    txt_path = jpg_path.replace('.jpg', '.txt')
    if os.path.exists(txt_path):
        with open(txt_path, 'r') as f:
            return f.read().strip()
    return ""

def write_output_file(file_path, paths):
    with open(file_path, 'w') as f:
        for jpg_path in paths:
            f.write(jpg_path + '\n')

def count_annotation_boxes(input_file):
    with open(input_file, 'r') as f:
        jpg_paths = [line.strip() for line in f.readlines()]
    
    num_with_annotation = 0
    num_without_annotation = 0

    for jpg_path in jpg_paths:
        txt_path = jpg_path.replace('.jpg', '.txt')
        if os.path.exists(txt_path) and os.path.getsize(txt_path) > 0:
            num_with_annotation += 1
        else:
            num_without_annotation += 1

    return num_with_annotation, num_without_annotation

# 使用示例
input_file = 'gt_guiji_all_20240909.txt'
output_files = ['gt_guiji_all_20240909_subset1.txt', 'gt_guiji_all_20240909_subset2.txt', 'gt_guiji_all_20240909_subset3.txt']
ratios = [1/3, 2/3, 1]  # 有标注框图片的分配比例
empty_ratios = [0, 0.1, 0.3]  # 无标注框图片的分配比例

num_with_annotation, num_without_annotation = count_annotation_boxes(input_file)

print("Number of images with annotation boxes:", num_with_annotation)
print("Number of images without annotation boxes:", num_without_annotation)

process_and_split_txt_file(input_file, output_files, ratios, empty_ratios)